import { BaseLoader } from '@cherrystudio/embedjs-interfaces'
import { cleanString, isValidURL } from '@cherrystudio/embedjs-utils'
import Embeddings from '@main/embeddings/Embeddings'
import { GlobalSemanticSplitter } from '@my_modules/lc-textsplitters/global-semantic-splitter.js'
//import { RecursiveCharacterTextSplitter } from '@langchain/textsplitters'
//import { RecursiveCharacterTextSplitter } from '@my_modules/lc-textsplitters'
import { SemanticSplitter } from '@my_modules/lc-textsplitters/semantic-splitter.js'
import { RecursiveCharacterTextSplitter } from '@my_modules/lc-textsplitters/text_splitter.js' // 引入 RecursiveCharacterTextSplitter
import { KnowledgeBaseParams } from '@types'
import md5 from 'md5'
import { getTextExtractor } from 'office-text-extractor'

export class PdfLoader extends BaseLoader<{ type: 'PdfLoader' }> {
  private readonly filePathOrUrl: string
  private readonly isUrl: boolean
  private readonly selectedStrategy: string
  private readonly semanticModel?: string
  private readonly semanticdimensions?: number
  private readonly semanticApiKey?: string
  private readonly semanticapiVersion?: string
  private readonly semanticapiBaseURL?: string

  constructor({
    filePathOrUrl,
    chunkOverlap,
    chunkSize,
    selectedStrategy,
    semanticModel,
    semanticdimensions,
    semanticApiKey,
    semanticapiVersion,
    semanticapiBaseURL
  }: {
    filePathOrUrl: string
    chunkSize?: number
    chunkOverlap?: number
    selectedStrategy: string
    semanticModel?: string
    semanticdimensions?: number
    semanticApiKey?: string
    semanticapiVersion?: string
    semanticapiBaseURL?: string
  }) {
    console.log('PdfLoader:', filePathOrUrl, chunkSize, chunkOverlap, selectedStrategy)
    super(
      `PdfLoader_${md5(filePathOrUrl)}`,
      { filePathOrUrl },
      chunkSize ?? 1000,
      chunkOverlap ?? 0,
      selectedStrategy ?? 'fixed'
    )

    this.filePathOrUrl = filePathOrUrl
    this.isUrl = isValidURL(filePathOrUrl) ? true : false
    this.selectedStrategy = selectedStrategy ?? 'fixed' // 添加这行
    this.semanticModel = semanticModel
    this.semanticdimensions = semanticdimensions
    this.semanticApiKey = semanticApiKey
    this.semanticapiVersion = semanticapiVersion
    this.semanticapiBaseURL = semanticapiBaseURL
  }

  // 生成分片器
  private _getSplitter() {
    switch (this.selectedStrategy) {
      case 'semantic': {
        const model = this.semanticModel
        const apiVersion = this.semanticapiVersion
        const baseURL = this.semanticapiBaseURL
        const dimensions = this.semanticdimensions
        const embeddings = new Embeddings({
          model,
          apiKey: this.semanticApiKey,
          apiVersion,
          baseURL,
          dimensions
        } as KnowledgeBaseParams)
        return new SemanticSplitter(embeddings, {
          minChunkSize: 200,
          maxChunkSize: 1000,
          similarityThreshold: 0.7,
          overlap: 0
        })
      }
      case 'overlap': {
        const model = this.semanticModel
        const apiVersion = this.semanticapiVersion
        const baseURL = this.semanticapiBaseURL
        const dimensions = this.semanticdimensions
        const embeddings = new Embeddings({
          model,
          apiKey: this.semanticApiKey,
          apiVersion,
          baseURL,
          dimensions
        } as KnowledgeBaseParams)
        return new GlobalSemanticSplitter(embeddings, {
          minChunkSize: 200,
          similarityThreshold: 0.7
        })
      }
      case 'fixed':
      default:
        return new RecursiveCharacterTextSplitter({
          chunkSize: this.chunkSize,
          chunkOverlap: this.chunkOverlap
        })
    }
  }

  override async *getUnfilteredChunks() {
    console.log(
      'strategy',
      this.selectedStrategy,
      this.semanticModel,
      this.semanticApiKey,
      this.semanticapiVersion,
      this.semanticapiBaseURL
    )

    const extractor = getTextExtractor()
    const pdfParsed = await extractor.extractText({ input: this.filePathOrUrl, type: this.isUrl ? 'url' : 'file' })

    console.log('pdfParsed:', pdfParsed)
    console.log('clean:', cleanString(pdfParsed))

    const splitter = this._getSplitter() // 使用新的私有方法获取分片器
    const chunks = await splitter.splitText(cleanString(pdfParsed))
    // const chunks = await semantic_split(cleanString(pdfParsed))
    console.log('totalchunk:', chunks)
    for (const chunk of chunks) {
      yield {
        pageContent: chunk,
        metadata: {
          type: 'PdfLoader' as const,
          source: this.filePathOrUrl
        }
      }
    }
  }
}
